{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import pybullet as p\n",
    "import pybullet_data\n",
    "from surrol.utils.pybullet_utils import (\n",
    "    step,\n",
    "    get_joints,\n",
    "    get_link_name,\n",
    "    reset_camera,\n",
    ")\n",
    "from surrol.robots.psm import Psm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling = 1.\n",
    "\n",
    "p.connect(p.GUI)\n",
    "# p.connect(p.DIRECT)\n",
    "p.setGravity(0, 0, -9.81)\n",
    "p.setAdditionalSearchPath(pybullet_data.getDataPath())\n",
    "reset_camera(yaw=90, pitch=-15, dist=0.9*scaling)\n",
    "\n",
    "# p.setPhysicsEngineParameter(contactBreakingThreshold=0.002)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "p.loadURDF(\"plane.urdf\", [0, 0, -0.001], globalScaling=1)\n",
    "\n",
    "psm = Psm((0, 0, 0.1524),\n",
    "          p.getQuaternionFromEuler((0, 0, -90/180*np.pi)), \n",
    "          scaling=scaling)\n",
    "psm.reset_joint((0, 0, 0.10, 0, 0, 0))\n",
    "\n",
    "# psm = Psm((0.05, 0.24, 0.8524),\n",
    "#           p.getQuaternionFromEuler((0, 0, np.deg2rad(-(90+20)))), \n",
    "#           scaling=scaling)\n",
    "# psm.reset_joint((0, 0, 0.10, 0, 0, 0))\n",
    "\n",
    "# for info in get_joints_info(psm.body, psm.joints):\n",
    "#     print(info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joints = get_joints(psm.body)\n",
    "print(\"There are {} joints.\\n\".format(len(joints)))\n",
    "\n",
    "for i in range(0, len(joints)):\n",
    "    print(get_link_name(psm.body, i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# continously run\n",
    "p.setRealTimeSimulation(1)\n",
    "\n",
    "while True:\n",
    "    p.setGravity(0, 0, -10)\n",
    "    time.sleep(0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forward Kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test with predefined pose\n",
    "# original joint position; 0.jpg\n",
    "\n",
    "psm.move_joint([-0.52359879, 0., 0.12, 0., 0., 0.])\n",
    "step(0.5)\n",
    "print(psm.get_current_position())\n",
    "\n",
    "# Read from dVRK (get_position_current)\n",
    "# [[-0.0077    0.8686   -0.4955   -0.0567]\n",
    "#  [ 0.9999    0.0001   -0.0154   -0.0002]\n",
    "#  [-0.0133   -0.4956   -0.8685   -0.0982]\n",
    "#  [      0         0         0    1.0000]]\n",
    "\n",
    "# previously compute\n",
    "# [[-6.12323415e-17  8.66025396e-01 -5.00000013e-01 -5.67500001e-02]\n",
    "#  [ 1.00000000e+00 -3.06161708e-17 -1.75493441e-16 -1.35258490e-17]\n",
    "#  [-1.67289863e-16 -5.00000013e-01 -8.66025396e-01 -9.82938802e-02]\n",
    "#  [ 0.00000000e+00  0.00000000e+00  0.00000000e+00  1.00000000e+00]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple joint position 2; 3.jpg\n",
    "\n",
    "psm.move_joint([-1.0471975511965976, 0.17453292519943295, 0.1, 0.7853981633974483, 0, 0])\n",
    "step(0.5)\n",
    "print(psm.get_current_position())\n",
    "\n",
    "# [[ 0.24721603  0.45989072 -0.85286855 -0.07974321]\n",
    "#  [ 0.69636423 -0.69636426 -0.17364817 -0.0162361 ]\n",
    "#  [-0.67376636 -0.55097853 -0.49240385 -0.04603976]\n",
    "#  [ 0.          0.          0.          1.        ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple joint position 3; 4.jpg\n",
    "\n",
    "psm.move_joint([-1.0471975511965976, 0.17453292519943295, 0.1, 0.7853981633974483, 0.5235987755982988, 0.2617993877991494])\n",
    "step(0.5)\n",
    "print(psm.get_current_position())\n",
    "\n",
    "# [[-0.21233893  0.66737769 -0.71380613 -0.07982825]\n",
    "#  [ 0.51624502 -0.54359788 -0.66180997 -0.01919286]\n",
    "#  [-0.82970071 -0.50902688 -0.22910342 -0.0423738 ]\n",
    "#  [ 0.          0.          0.          1.        ]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open/Close Jaw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open jaw test\n",
    "\n",
    "psm.open_jaw()\n",
    "step(0.5)\n",
    "print(psm.get_current_position())\n",
    "print(np.rad2deg(psm.get_current_jaw_position()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# close jaw test\n",
    "\n",
    "psm.close_jaw()\n",
    "step(0.5)\n",
    "print(np.rad2deg(psm.get_current_jaw_position()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inverse Kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inverse kinematics test\n",
    "# simple coord position 0; 5.jpg\n",
    "pose = np.array([\n",
    "    [ 4.33293348e-13,  1.00000000e+00,  4.52526905e-13, -1.33425265e-17],\n",
    "    [ 1.00000000e+00, -4.33293348e-13,  9.95143306e-14, -1.12816941e-17],\n",
    "    [ 9.95143306e-14,  4.52526905e-13, -1.00000000e+00, -1.13499997e-01],\n",
    "    [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  1.00000000e+00]\n",
    "])\n",
    "psm.move(pose)\n",
    "step(0.5)\n",
    "print(psm.get_current_joint_position())\n",
    "print(psm.get_current_position())\n",
    "\n",
    "# [ 0.00000000e+00  8.65586272e-15  1.19999997e-01  4.33293348e-13\n",
    "#  -1.08357767e-13  4.52504218e-13]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple coord position 1; 6.jpg\n",
    "pose = np.array([\n",
    "    [ 4.33293348e-13,  1.00000000e+00,  4.52526905e-13,  5.00000000e-02],\n",
    "    [ 1.00000000e+00, -4.33293348e-13,  9.95143306e-14, -1.12816941e-17],\n",
    "    [ 9.95143306e-14,  4.52526905e-13, -1.00000000e+00, -1.13499997e-01],\n",
    "    [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  1.00000000e+00]\n",
    "])\n",
    "psm.move(pose)\n",
    "step(0.5)\n",
    "print(psm.get_current_joint_position())\n",
    "print(psm.get_current_position())\n",
    "\n",
    "# [ 4.14949688e-01 -7.10542736e-15  1.30525197e-01  4.36665965e-13\n",
    "#   9.00240590e-14 -4.14949688e-01]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JIGSAWS Kinematics Record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# jigsaws test\n",
    "# import kinematics data transformed from JIGSAWS \n",
    "# (https://cirl.lcsr.jhu.edu/research/hmm/datasets/jigsaws_release/)\n",
    "\n",
    "joint_values = np.load('qs_jigsaws.npy')\n",
    "psm.close_jaw()\n",
    "step(1)\n",
    "\n",
    "start_time = time.time()\n",
    "for i in tqdm(range(len(joint_values))):\n",
    "    psm.move_joint(joint_values[i])\n",
    "    psm.close_jaw()\n",
    "    step(0.5)\n",
    "    _ = p.getCameraImage(128, 128)\n",
    "\n",
    "end_time = time.time()\n",
    "print(\"Used time: {:.4f}\".format(end_time - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = p.getCameraImage(512, 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "surrol",
   "language": "python",
   "name": "rorl"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
